package com.google.code.b0rx0r.ui;

/**
 * This code is released into the public domain
 */

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.swing.event.TreeModelEvent;
import javax.swing.event.TreeModelListener;
import javax.swing.tree.TreeModel;
import javax.swing.tree.TreePath;

/**
 * @author Rogan Dawes <SortedTreeModel @ dawes.za.net>
 */
public class SortedTreeModel implements TreeModel {

	private TreeModel delegate;

	private Comparator<Object> comparator = null;

	private Map<Object, int[]> viewToModel = new HashMap<Object, int[]>();

	private Listener listener = new Listener();

	/**
	 * Constructs a sorted TreeModel, based on the data in the provided delegate,
	 * sorted according to the initial Comparator provided
	 * 
	 * @param delegate
	 *          the underlying TreeModel to wrap
	 * @param comparator
	 *          the initial Comparator to use when sorting the nodes
	 */
	public SortedTreeModel(TreeModel delegate, Comparator<Object> comparator) {
		this.delegate = delegate;
		setComparator(comparator);
		delegate.addTreeModelListener(listener);
	}

	public Object getChild(Object parent, int index) {
		int[] viewToModel = getViewToModel(parent);
		if (viewToModel == null)
			return delegate.getChild(parent, index);
		return delegate.getChild(parent, viewToModel[index]);
	}

	public int getChildCount(Object parent) {
		return delegate.getChildCount(parent);
	}

	public int getIndexOfChild(Object parent, Object child) {
		int index = delegate.getIndexOfChild(parent, child);
		int[] viewToModel = getViewToModel(parent);
		if (viewToModel == null)
			return index;
		for (int i = 0; i < viewToModel.length; i++)
			if (viewToModel[i] == index)
				return i;
		throw new RuntimeException("This should never happen"); //$NON-NLS-1$
	}

	public Object getRoot() {
		return delegate.getRoot();
	}

	public boolean isLeaf(Object node) {
		return delegate.isLeaf(node);
	}

	public void valueForPathChanged(TreePath path, Object newValue) {
		delegate.valueForPathChanged(path, newValue);
	}

	/**
	 * Set the new Comparator to use to sort the underlying tree nodes
	 * 
	 * @param comparator
	 *          the comparator
	 */
	public void setComparator(Comparator<Object> comparator) {
		this.comparator = comparator;
		if (sort(new TreePath(getRoot()), true))
			fireStructureChanged();
	}

	/**
	 * Sorts the children of the node at the specified path, and optionally all of
	 * its children
	 * 
	 * @param path
	 *          the path to the node to sort
	 * @param recursive
	 *          whether to sort that node's children as well
	 * @return true if any changes were made to the order of the nodes, false
	 *         otherwise
	 */
	protected boolean sort(TreePath path, boolean recursive) {
		Object parent = path.getLastPathComponent();
		int childCount = delegate.getChildCount(parent);
		if (childCount == 0)
			return false;
		Object[] children = new Object[childCount];
		int[] viewToModel = new int[childCount];
		for (int i = 0; i < childCount; i++) {
			children[i] = delegate.getChild(parent, i);
		}
		Arrays.sort(children, comparator);
		for (int i = 0; i < childCount; i++)
			viewToModel[i] = delegate.getIndexOfChild(parent, children[i]);
		boolean changed = setViewToModel(parent, viewToModel);
		if (recursive)
			for (int i = 0; i < childCount; i++)
				changed |= sort(path.pathByAddingChild(children[i]), recursive);
		return changed;
	}

	private int[] getViewToModel(Object parent) {
		return viewToModel.get(parent);
	}

	/**
	 * Store the array mapping entries from the underlying model to the sorted
	 * order As an optimization, if the sort order is identical to the underlying
	 * data order, we do not store a mapping
	 * 
	 * @param parent
	 *          the node whose children are represented
	 * @param mapping
	 *          the mapping from view order to underlying model order
	 * @return true if the mapping differs from the current mapping, false
	 *         otherwise
	 */
	private boolean setViewToModel(Object parent, int[] mapping) {
		boolean identity = true;
		for (int i = 0; i < mapping.length; i++)
			if (mapping[i] != i) {
				identity = false;
				break;
			}
		if (identity)
			mapping = null;
		int[] oldMapping = getViewToModel(parent);
		if (!equals(oldMapping, mapping)) {
			if (mapping == null) {
				viewToModel.remove(parent);
			} else {
				viewToModel.put(parent, mapping);
			}
			return true;
		}
		return false;
	}

	/**
	 * Compare two integer arrays for equality
	 * 
	 * @param a
	 *          an array
	 * @param b
	 *          an array
	 * @return true if both arrays are null, or contain the same values in order
	 */
	private boolean equals(int[] a, int[] b) {
		if (a == b)
			return true;
		if (a == null || b == null)
			return false;
		if (a.length != b.length)
			return false;
		for (int i = 0; i < a.length; i++)
			if (a[i] != b[i])
				return false;
		return true;
	}

	private class Listener implements TreeModelListener {

		/**
		 * remaps the indices provided in the event to those in the sorted model
		 * 
		 * @param path
		 *          the path of the parent node
		 * @param indices
		 *          the unsorted indices
		 */
		private int[] remapIndices(TreePath path, int[] indices) {
			int[] mapping = getViewToModel(path.getLastPathComponent());
			if (mapping == null)
				return indices;

			int[] newIndices = new int[indices.length];
			for (int i = 0; i < indices.length; i++) {
				for (int j = 0; j < mapping.length; j++) {
					if (mapping[j] == indices[i]) {
						newIndices[i] = j;
						break;
					}
				}
			}
			return newIndices;
		}

		/*
		 * (non-Javadoc)
		 * 
		 * @see javax.swing.event.TreeModelListener#treeNodesChanged(javax.swing.event.TreeModelEvent)
		 */
		public void treeNodesChanged(TreeModelEvent e) {
			int[] indices = e.getChildIndices();
			TreePath path = e.getTreePath();
			sort(path, true);
			indices = remapIndices(path, indices);
			fireTreeStructureChanged(path);
		}

		/*
		 * (non-Javadoc)
		 * 
		 * @see javax.swing.event.TreeModelListener#treeNodesInserted(javax.swing.event.TreeModelEvent)
		 */
		public void treeNodesInserted(TreeModelEvent e) {
			int[] indices = e.getChildIndices();
			TreePath path = e.getTreePath();
			sort(e.getTreePath(), false);
			indices = remapIndices(path, indices);
			fireChildrenAdded(path, indices, e.getChildren());
		}


		/*
		 * (non-Javadoc)
		 * 
		 * @see javax.swing.event.TreeModelListener#treeNodesRemoved(javax.swing.event.TreeModelEvent)
		 */
		public void treeNodesRemoved(TreeModelEvent e) {
			int[] indices = e.getChildIndices();
			TreePath path = e.getTreePath();
			indices = remapIndices(path, indices);
			sort(e.getTreePath(), false);
			fireChildrenRemoved(path, indices, e.getChildren());
		}

		/*
		 * (non-Javadoc)
		 * 
		 * @see javax.swing.event.TreeModelListener#treeStructureChanged(javax.swing.event.TreeModelEvent)
		 */
		public void treeStructureChanged(TreeModelEvent e) {
			fireTreeStructureChanged(e.getTreePath());
		}


	}

	public static void main(String[] args) {
		final javax.swing.tree.DefaultMutableTreeNode root = new javax.swing.tree.DefaultMutableTreeNode(Messages.getString("SortedTreeModel.1")); //$NON-NLS-1$
		buildTree(root, 5, 3);
		final javax.swing.tree.DefaultTreeModel delegate = new javax.swing.tree.DefaultTreeModel(root);
		final SortedTreeModel sorted = new SortedTreeModel(delegate, new StringNodeComparator(1, true));
		final javax.swing.JTree tree = new javax.swing.JTree(sorted);
		javax.swing.JFrame frame = new javax.swing.JFrame(Messages.getString("SortedTreeModel.2")); //$NON-NLS-1$
		javax.swing.JScrollPane scrollPane = new javax.swing.JScrollPane(tree);
		javax.swing.JButton sort = new javax.swing.JButton(Messages.getString("SortedTreeModel.3")); //$NON-NLS-1$
		final StringNodeComparator[] comparators = new StringNodeComparator[] { new StringNodeComparator(-1, false), new StringNodeComparator(1, false), new StringNodeComparator(-1, true), new StringNodeComparator(1, true), };
		sort.addActionListener(new java.awt.event.ActionListener() {
			private int i = 0;

			private java.util.List<TreePath> expansion = new java.util.ArrayList<TreePath>();

			private TreePath[] selection;

			public void actionPerformed(java.awt.event.ActionEvent ae) {
				saveSelectionAndExpansion();

				System.out.println(comparators[i]);
				sorted.setComparator(comparators[i]);
				i = (i + 1) % comparators.length;

				restoreSelectionAndExpansion();
			}

			private void saveSelectionAndExpansion() {
				expansion.clear();
				java.util.Enumeration<TreePath> e = tree.getExpandedDescendants(new TreePath(sorted.getRoot()));
				while (e.hasMoreElements())
					expansion.add(e.nextElement());
				selection = tree.getSelectionPaths();
			}

			private void restoreSelectionAndExpansion() {
				tree.setSelectionPaths(selection);
				java.util.Iterator<TreePath> it = expansion.iterator();
				while (it.hasNext())
					tree.expandPath(it.next());
			}
		});
		javax.swing.JButton mutate = new javax.swing.JButton(Messages.getString("SortedTreeModel.4")); //$NON-NLS-1$
		mutate.addActionListener(new java.awt.event.ActionListener() {
			private boolean running = false;

			private Thread mutator = new Thread(new Runnable() {
				public void run() {
					try {
						java.util.Random random = new java.util.Random();
						while (running) {
							int i = Math.abs(random.nextInt()) % delegate.getChildCount(root);
							final javax.swing.tree.MutableTreeNode child = (javax.swing.tree.MutableTreeNode) delegate.getChild(root, i);
							java.awt.EventQueue.invokeLater(new Runnable() {
								public void run() {
									delegate.removeNodeFromParent(child);

								}
							});
							Thread.sleep(1000);
							java.awt.EventQueue.invokeLater(new Runnable() {
								public void run() {
									delegate.insertNodeInto(child, root, 0);
								}
							});
							Thread.sleep(1000);
						}
					} catch (InterruptedException ie) {
					}
				}
			});

			public void actionPerformed(java.awt.event.ActionEvent ae) {
				if (!running) {
					running = true;
					mutator.start();
				} else {
					running = false;
				}
			}
		});
		frame.setDefaultCloseOperation(javax.swing.JFrame.EXIT_ON_CLOSE);
		frame.getContentPane().setLayout(new java.awt.BorderLayout());
		frame.getContentPane().add(scrollPane, java.awt.BorderLayout.CENTER);
		frame.getContentPane().add(sort, java.awt.BorderLayout.NORTH);
		frame.getContentPane().add(mutate, java.awt.BorderLayout.SOUTH);
		frame.setBounds(200, 200, 400, 400);
		frame.setVisible(true);
	}

	private static void buildTree(javax.swing.tree.MutableTreeNode parent, int breadth, int depth) {
		if (depth == 0)
			return;
		for (int i = 0; i < breadth; i++) {
			javax.swing.tree.MutableTreeNode node = new javax.swing.tree.DefaultMutableTreeNode(parent + Messages.getString("SortedTreeModel.5") + i); //$NON-NLS-1$
			parent.insert(node, i);
			buildTree(node, breadth, depth - 1);
		}
	}

	private static class StringNodeComparator implements Comparator<Object> {

		private int order;

		private boolean leavesOnly;

		public StringNodeComparator(int order, boolean leavesOnly) {
			this.order = order;
			this.leavesOnly = leavesOnly;
		}

		/*
		 * (non-Javadoc)
		 * 
		 * @see java.util.Comparator#compare(java.lang.Object, java.lang.Object)
		 */
		public int compare(Object o1, Object o2) {
			javax.swing.tree.TreeNode t1 = (javax.swing.tree.TreeNode) o1;
			javax.swing.tree.TreeNode t2 = (javax.swing.tree.TreeNode) o2;
			int c = o1.toString().compareTo(o2.toString());
			if (leavesOnly)
				if (t1.isLeaf() && t2.isLeaf()) {
					return order * c;
				} else
					return c;
			return order * c;
		}

		public String toString() {
			return Messages.getString("SortedTreeModel.6") + (leavesOnly ? Messages.getString("SortedTreeModel.7") : Messages.getString("SortedTreeModel.8")) + Messages.getString("SortedTreeModel.9") + (order == -1 ? Messages.getString("SortedTreeModel.10") : Messages.getString("SortedTreeModel.11")) + Messages.getString("SortedTreeModel.12"); //$NON-NLS-1$ //$NON-NLS-2$ //$NON-NLS-3$ //$NON-NLS-4$ //$NON-NLS-5$ //$NON-NLS-6$ //$NON-NLS-7$
		}
	}
	
	private List<TreeModelListener> listeners = new ArrayList<TreeModelListener>();

	@Override
	public void addTreeModelListener(TreeModelListener l) {
		listeners.add(l);
	}

	@Override
	public void removeTreeModelListener(TreeModelListener l) {
    listeners.remove(l);
	}
	

	private void fireStructureChanged() {
		fireTreeStructureChanged(new TreePath(delegate.getRoot()));
	}

	private void fireChildrenChanged(TreePath path, int[] indices, Object[] children) {
    TreeModelEvent e = new TreeModelEvent(this, path, indices, children);
		for (TreeModelListener l : listeners) {
			l.treeNodesChanged(e);
		}
	}

	private void fireChildrenAdded(TreePath path, int[] indices, Object[] children) {
    TreeModelEvent e = new TreeModelEvent(this, path, indices, children);
		for (TreeModelListener l : listeners) {
			l.treeNodesInserted(e);
		}
	}

	private void fireChildrenRemoved(TreePath path, int[] indices, Object[] children) {
    TreeModelEvent e = new TreeModelEvent(this, path, indices, children);
		for (TreeModelListener l : listeners) {
			l.treeNodesRemoved(e);
		}
	}

	private void fireTreeStructureChanged(TreePath treePath) {
		TreeModelEvent e = new TreeModelEvent(this, treePath);
		for (TreeModelListener l : listeners) {
			l.treeStructureChanged(e);
		}
	}
}